import torch
import torch.nn as nn


# model
# dataloader
# loss

class Conv_Bn_Relu(nn.Module):
    def __init__(self, c1, c2, k=3, s=1, p=0):
        super(Conv_Bn_Relu, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, p)
        self.bn = nn.BatchNorm2d(c2)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))


class AlexNet(nn.Module):
    def __init__(self, num_classes, init_weight=False):
        """
        输入尺寸为 224 * 224
        Args:
            num_classes: 分类数
            init_weight: 是否初始化
        """
        super(AlexNet, self).__init__()
        self.feature = nn.Sequential(
            Conv_Bn_Relu(3, 48, 11, 4, 2),
            nn.MaxPool2d(kernel_size=3, stride=2),
            Conv_Bn_Relu(48, 128, 5, 2),
            nn.MaxPool2d(kernel_size=3, stride=2),
            Conv_Bn_Relu(128, 192, 3, 1, 1),
            Conv_Bn_Relu(192, 192, 3, 1, 1),
            Conv_Bn_Relu(192, 128, 3, 1, 1),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )

        self.classifier = nn.Sequential(  # 打包模块
            nn.Dropout(p=0.5),
            # 全链接
            nn.Linear(128 * 2 * 2, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),  # 防止过拟合
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )

        if init_weight:
            self._initialize_weights()

    def forward(self, x):
        x = self.feature(x)
        x = torch.flatten(x, start_dim=1)  # 上一步 x [1, 128, 2, 2], 需要展开成 [1, 128 * 2 * 2] 才能送入全连接层
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')  # 何教授方法
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)  # 正态分布赋值
                nn.init.constant_(m.bias, 0)


if __name__ == '__main__':
    x = torch.rand([1, 3, 224, 224])
    model = AlexNet(num_classes=10, init_weight=False)
    y = model(x)
    print(y)
